from pyspark.sql import DataFrame
import pyspark.sql.functions as F
from pyspark.sql.types import StringType

from tag.BaseModel import BaseModel


class NationalityTag(BaseModel):
    def compute(self, es_df, five_df):
        """
        重写父类中的compute方法，用于国籍标签计算
        :param es_df: 业务数据
        :param five_df: 标签规则数据
        :return: 新的标签结果
        """
        # 打印输入数据结构和内容（调试用）
        es_df.printSchema()
        es_df.show()
        five_df.printSchema()
        five_df.show()
        # 1. 执行标签匹配（左连接保留所有用户）
        # 注意：使用es_df的'nationality'字段而不是'job'字段
        new_df = es_df.join(
            other=five_df,
            on=es_df['nationality'] == five_df['rule'],
            how='left'
        ).select(es_df['id'].alias("userId"), five_df['id'].alias("tagsId"))

        return new_df


if __name__ == '__main__':
    # 假设国籍标签的四级ID是70（需替换为实际值）
    nationalityModel = NationalityTag(70)
    nationalityModel.execute()